#include <string.h>
#include <stdlib.h>
#include <ctype.h>
#include <stdio.h>
#include <mpi.h>
#include <omp.h>
#include "avl.h"

#define MAX_WORDS       10000000
#define MAX_PRIME       8000
#define MAX_WORD_SIZE   300

void error(const char *message);
void divide(int nprocesses, const char *filename);
void conquer(int rank);

int main(int argc, char *argv[])
{
    /* check for file argument */
    if(argc < 2)
    {
        fprintf(stderr, "Missing input file.\n");
        return 1;
    }

    /* initialize MPI */
    if(MPI_Init(&argc, &argv) == MPI_SUCCESS)
    {
        int nprocesses, rank;

        MPI_Comm_size(MPI_COMM_WORLD, &nprocesses);

        /* we need at least two processes */
        if(nprocesses < 2) error("Set at least two processes.");

        MPI_Comm_rank(MPI_COMM_WORLD, &rank);

        /* the root process (rank 0) will divide the problem and send minor parts of it
         * to the others processes */
        if(rank == 0)
            divide(nprocesses, argv[1]);
        else
            conquer(rank);
    }
    else
    {
        fprintf(stderr, "MPI Error.\n");
        return 1;
    }

    MPI_Finalize();

    return 0;
}

void error(const char *message)
{
    fprintf(stderr, "%s\n", message);
    MPI_Abort(MPI_COMM_WORLD, 1);
}

void divide(int nprocesses, const char *filename)
{
    fprintf(stderr, "root: Opening file...\n");

    /* open file and check if it exists */
    FILE *fp = fopen(filename, "r");
    if(fp == NULL) error("File does not exists.");

    /* retrieve file size */
    fseek(fp, 0, SEEK_END);
    long filesize = ftell(fp);

    /* allocate space to store file */
    char *buffer = (char *) malloc(filesize * sizeof(char));

    if(buffer == NULL) error("Out of memory.");

    /* copy file to memory and close it */
    rewind(fp);
    fread(buffer, sizeof(char), filesize, fp);
    fclose(fp);

    /* simulate (or try to simulate) a three stage pipeline:
     * 1st stage: substitute symbols by spaces and turn letters lowercase
     * 2nd stage: look for words beginnings and ends
     * 3rd stage: cut repeated words and send them
     */

    /* control vars */
    int idx1 = 0; /* indicate in which point first stage is */
    int idx2 = 0; /* indicate in which point second stage is */
    int idx3 = 0; /* indicate in which point third stage is */

    /* allocate space to storage found words */
    char **words = (char **) malloc((MAX_WORDS + 1) * sizeof(char *));
    size_t *sizes = (size_t *) malloc(MAX_WORDS * sizeof(size_t));

    if(words == NULL || sizes == NULL) error("Out of memory.");

    fprintf(stderr, "root: Sending words...\n");
    #pragma omp parallel shared(buffer,idx1,idx2,idx3,words,sizes) \
                         firstprivate(nprocesses) num_threads(3)
    {
        #pragma omp sections
        {
            /* 1 */
            #pragma omp section
            {
                for(idx1 = 0; idx1 < filesize; idx1++)
                    if(isalnum(buffer[idx1])) /* is letter or digit */
                        buffer[idx1] = tolower(buffer[idx1]);
                    else
                        buffer[idx1] = ' ';
                /* sum a big number in order to unlock the second stage */
                idx1 += 5 * MAX_WORD_SIZE;
            }
            /* 2 */
            #pragma omp section
            {
                /* wait for a while */
                while(10 * MAX_WORD_SIZE > idx1);

                char *pch = strtok(buffer, " ");

                for(idx2 = 0; idx2 < MAX_WORDS && pch != NULL; idx2++)
                {
                    words[idx2] = pch;
                    sizes[idx2] = strlen(pch) + 1;

                    /* sync to stage 1 */
                    while(pch - buffer + 2 * MAX_WORD_SIZE > idx1);

                    pch = strtok(NULL, " ");
                }

                words[idx2++] = NULL;

                if(idx2 == MAX_WORDS + 1)
                    fprintf(stderr, "Maximum of words reached, consider increasing MAX_WORDS\n");
                else
                    fprintf(stderr, "root: %d words were found.\n", idx2);
            }
            /* 3 */
            #pragma omp section
            {
                /* sync to stage 2 */
                while(idx2 <= idx3);

                avl_t *avl = avl_create();
                int process = 0;

                for(idx3 = 0; words[idx3] != NULL; idx3++)
                {
                    /* verify if the word is repeated */
                    if(sizes[idx3] >= 3 && avl_insert(avl, words[idx3]))
                    {
                        /* send word to a process */
                        MPI_Send(words[idx3], sizes[idx3], MPI_CHAR,
                                  process + 1, 0, MPI_COMM_WORLD);
                        process = (process + 1) % (nprocesses - 1);
                    }
                    /* sync to stage 2 */
                    while(idx2 <= idx3);
                }

                avl_destroy(avl);

                /* empty string indicates the end of the process */
                for(process = 1; process < nprocesses; process++)
                    MPI_Send("", 1, MPI_CHAR, process, 0, MPI_COMM_WORLD);
            }
        }
    }
    /* end parallel section */

    fprintf(stderr, "root: Finished!\n");

    free(sizes);
    free(words);
    free(buffer);
}

int ispalindrome(const char *word, int len);

void conquer(int rank)
{
    fprintf(stderr, "%d: Calculating Sieve of Eratosthenes...\n", rank);

    /* the vector isnotprime has 1 if the number is NOT prime and 0 otherwise */
    int *isnotprime = (int *) calloc(MAX_PRIME + 1, sizeof(int));

    if(isnotprime == NULL) error("Out of memory.");

    /* calculate Sieve of Eratosthenes up to MAX_PRIME */
    isnotprime[0] = 1;
    isnotprime[1] = 1;

    int i, j;
    for (i = 2; i <= MAX_PRIME; i++) {
        if(isnotprime[i]) continue; /* it was already marked */
        /* mark multiples */
        for(j = i * i; j <= MAX_PRIME ; j += i)
            isnotprime[j] = 1;
    }

    /* receive word from root process */
    MPI_Status status;
    char word[MAX_WORD_SIZE];
    int wordcount = 0;
    int len;

    fprintf(stderr, "%d: Receiving messages from root and processing...\n", rank);

    MPI_Recv(word, MAX_WORD_SIZE, MPI_CHAR, 0, 0, MPI_COMM_WORLD, &status);
    MPI_Get_count(&status, MPI_CHAR, &len);
    len--; /* ignore \0 */

    /* stop when it receives a empty string */
    while(len > 0)
    {
        wordcount++;
        /* start word processing */
        if(ispalindrome(word, len))
        {
            int k, sum = 0;
            for(k = 0; k < len; k++)
                sum += (int) word[k];

            printf("%s - %d %s prime\n", word, sum,
                   isnotprime[sum] ? "is not" : "is");
        }

        /* receive next word */
        MPI_Recv(word, MAX_WORD_SIZE, MPI_CHAR, 0, 0, MPI_COMM_WORLD, &status);
        MPI_Get_count(&status, MPI_CHAR, &len);
        len--;
    }

    fprintf(stderr, "%d: %d words were received and processed\n", rank, wordcount);
    fprintf(stderr, "%d: Finished!\n", rank);

    free(isnotprime);
}

int ispalindrome(const char *word, int len)
{
    int i, j;
    for(i = 0, j = len - 1; i <= j; i++, j--)
        if(word[i] != word[j])
            return 0;
    return 1;
}

